//============================================================================
// Name        : Coevolution 3
// Author      : Alex Seeholzer
// Description :
// Extension of Coevolution 2. Adding several random viral Wildtypes with higher fitness, to guide
// virus mutation in the functional region.
//
// Usage: ./QuasiSpecies .0003 .0001 5. 5. .3 15. 25
//============================================================================
// cin with strings
#include <iostream>
#include <string>
#include <stdio.h>
#include <math.h>
#include <typeinfo>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_interp.h>
using namespace std;

#define TYPE unsigned int
#define L 32
#define ONE 1
#define RRMASK 0x000000FF
#define MAXINT 0xFFFFFFFF
#define LRR 8
#define VIR 0
#define IS 1
#define DELTA 1.
#define N 150000
#define TMAX 1000

#define NM 2 /* population number, 0=virus 1=immune */
const TYPE masters[] = {0x88888800, 0xFFFFFF00}; /* the individual master sequences for structure */
bool vir_masters[(ONE<<LRR)]; // array of master sequences vor virus functional region, true == master

double alpha[2];
double base;
double beta;
double gam[2];
double initratio;
int tmax;

inline TYPE mutate (TYPE x, double mu);
inline int hd (TYPE x, TYPE y); // hamming dist
inline void outbinary(TYPE num);
void reproduce (TYPE pop[NM][N], double* mu, unsigned int cdis[NM][(ONE << LRR)]);
void sample(TYPE* pop,TYPE* popnew,double* mut,double* sel,double mu);
TYPE compute_consensus (TYPE *pop);

// computes for a pop the distribution of clones (stored in cdis)
void compute_cdis (TYPE *pop, unsigned int *cdis, TYPE master);

// Random number generator
gsl_rng *r;
gsl_interp_accel *acc;
const gsl_rng_type *T;

int main(int argc, char **argv) {

	// Set up output file
	FILE * outfile;
	int i,j,k,t,s,vir_masters_num;
	TYPE pop[NM][N]; // array of population
	double mu[NM]; // mutation rate
	TYPE smax[NM]; // maximal populated sequences for each population
	unsigned int cdis[NM][(ONE << LRR)]; // for each population, distribution of functional region

	// random number generators
	gsl_rng_env_setup();
	acc = gsl_interp_accel_alloc();
	T = gsl_rng_gfsr4;
	r = gsl_rng_alloc(T);

	// read mu, s = seed
	if (argc > 12) {
		mu[VIR]=atof(argv[1]);
		mu[IS]=atof(argv[2]);
		base=atof(argv[3]);
		alpha[VIR]=atof(argv[4]);
		alpha[IS]=atof(argv[5]);
		gam[VIR]=atof(argv[6]);
		gam[IS]=atof(argv[7]);
		vir_masters_num=atoi(argv[8]);
		beta=atof(argv[9]);
		initratio=atof(argv[10]);
		s=atoi(argv[11]);
		tmax=atoi(argv[12]);
		fprintf(stderr, "\nmu_v=% .3e, mu_is=% .3e, base=% .3e, alpha_v=% .3e, alpha_is=% .3e, gamma_v=% .3e, gamma_is=% .3e, peaks=%i, beta=% .3e, initratio=%.3e, N= %i, s=%i, tmax=%i\n", mu[VIR], mu[IS], base, alpha[VIR], alpha[IS], gam[VIR], gam[IS], vir_masters_num, beta, initratio, N, s, tmax);
	} else {
		fprintf(stderr, "not enough args. argc=%i\n", argc);
		exit(1);
	}
	char *outfilename = (char*) malloc(255);
	sprintf(outfilename, "run_%s_%s_%s_%s_%s_%s_%s_%s_%s_%s_%s_%s_%i.dat", argv[1],argv[2],argv[3],argv[4],argv[5],argv[6],argv[7],argv[8],argv[9],argv[10],argv[11],argv[12],N);
	outfile = fopen(outfilename,"w");
	free(outfilename);
	fprintf(outfile, "mu_v=% .3e, mu_{is}=% .3e, base=% .3e, alpha_v=% .3e, alpha_{is}=% .3e, gamma_v=% .3e, gamma_{is}=% .3e, peaks=%i, beta=% .3e, initratio=%.3e, N= %i, s=%i, tmax=%i\n", mu[VIR], mu[IS], base, alpha[VIR], alpha[IS], gam[VIR], gam[IS], vir_masters_num, beta, initratio, N, s, tmax);
	gsl_rng_set (r, s); // initialize random num generator with seed s

	// compute vir_masters_num unique fitness peaks in [0,(ONE<<LRR)] and store in vir_masters
	int vir_masters_idx[vir_masters_num];
	TYPE b[(ONE<<LRR)];
	for (i = 0; i < (ONE<<LRR); i++) {
		b[i] = (TYPE) i;
	}
	gsl_ran_choose (r, vir_masters_idx, vir_masters_num, b, (ONE<<LRR), sizeof (TYPE));
	for(i=0;i<(ONE<<LRR);i++)
		vir_masters[i]=false;
	for(i=0;i<vir_masters_num;i++)
		vir_masters[vir_masters_idx[i]]=true;

	unsigned int ham = hd(vir_masters_idx[0],vir_masters_idx[1]);
	fprintf(stderr,"\nMaster sequences for virus population: (hamdist = %i)\n",ham);
	fprintf(outfile,"Master sequences for virus population: (hamdist = %i)\n",ham);
	for(i=0;i<vir_masters_num;i++) {
		fprintf(stderr,"%i ",vir_masters_idx[i]);
		fprintf(outfile,"%i ",vir_masters_idx[i]);
	}
	fprintf(outfile,"\n");
	fprintf(stderr,"\n");
	fprintf(stderr,"\n");

	// immune population initialization, set as master + random sequence initially
	for (i = 0; i < N; i++)
		pop[IS][i] = (masters[IS] & ~RRMASK)
				| gsl_rng_uniform_int(r, ONE << LRR);

	// virus population initialization, set as master + random sequence initially
	for (i = 0; i < N; i++) {
		int mstr = (i>=N*initratio) ?  vir_masters_idx[0] : vir_masters_idx[1];
		pop[VIR][i] = (masters[VIR] & ~RRMASK)
						| mstr;
	}

	// compute frequency of master sequences for initial reproduce call
	for (j = 0; j < NM; j++)
		compute_cdis(pop[j], cdis[j], masters[j]);

	// main loop - reproduce, calc distribution and output
	int mover = 2;
	for (t = 0; t < tmax; t++) {

		// manual shift of virus population for tryout
		if (false && t==500) {
			for (i = 0; i < N; i++)
				pop[VIR][i] = (masters[VIR] & ~RRMASK)
						| vir_masters_idx[mover%vir_masters_num];
			mover++;
		}

		reproduce(pop, mu, cdis);
		for (j = 0; j < NM; j++)
			compute_cdis(pop[j], cdis[j], masters[j]);

		for (k = 0; k < (ONE << LRR); k++) {
			fprintf(outfile,"%i %i ", t, k);
			for (j = 0; j < NM; j++)
				fprintf(outfile,"%i ", cdis[j][k]);
			fprintf(outfile,"\n");
		}
		if (t%100 == 0)
			fprintf(stderr,"\r%.0f%%",t/(double)tmax*100.);
	}

	fprintf(stderr,"\n");
	fprintf(stderr,"\n");

	for (j = 0; j < NM; j++) {
		fprintf(stderr,"Population %i\n",j);
		smax[j] = 0;
		for (k = 1; k < (ONE << LRR); k++)
			if (cdis[j][k] > cdis[j][smax[j]])
				smax[j] = k;
		fprintf(stderr, "smax[%i]= %3i (%i), neighbors: ", j, smax[j],
				cdis[j][smax[j]]);
		for (k = 0; k < LRR; k++)
			fprintf(stderr, "%3i (%i), ", smax[j] ^ (ONE << k), cdis[j][(smax[j] ^ (ONE << k))]);
		fprintf(stderr, "\n");
	}

	fclose(outfile);
	return 0;
}

void compute_cdis (TYPE *pop, unsigned int *cdis, TYPE master) {
  for (int k=0; k<(ONE << LRR); k++)
    cdis[k]=0;

  for (int i=0; i<N; i++)
    if ((pop[i] & ~RRMASK)==(master & ~RRMASK))
      cdis[(unsigned int)(pop[i] & RRMASK)]++;

  return;
}

void reproduce(TYPE pop[NM][N], double* mu, unsigned int cdis[NM][(ONE << LRR)]) {
	int i, j;
	TYPE popnew[NM][N];
	double sel[N];
	double mut[N];
	double delta;

	// reproduce virus population
	for (i = 0; i < N; i++) {
		// calculate percentage of genome pop[VIR][i] in pop[IS]
		delta = cdis[IS][(pop[VIR][i]) & RRMASK] / (double) N;
		mut[i] = base;
		sel[i] = (1/DELTA)-(gam[VIR]*delta);
		if(sel[i]<0) {
			fprintf(stderr,"<0");
			sel[i]=0;
		}
		if ((pop[VIR][i] & ~RRMASK) == (masters[VIR] & ~RRMASK)) // if we have master sequence
			mut[i] = alpha[VIR];
		if (vir_masters[pop[VIR][i] & RRMASK]) // if we have one of the functional master sequences
			mut[i] += beta;
	}
	sample(pop[VIR],popnew[VIR],mut,sel,mu[VIR]);

	// reproduce immune system population
	for (i = 0; i < N; i++) {
		// calculate percentage of genome pop[IS][i] in pop[VIR]
		delta = cdis[VIR][(pop[IS][i]) & RRMASK] / (double) N;
		mut[i] = base;
		sel[i] = (1/DELTA);
		if ((pop[IS][i] & ~RRMASK) == (masters[IS] & ~RRMASK)) // if we have master sequence
			mut[i] = alpha[IS]+(gam[IS]*delta);
	}
	sample(pop[IS],popnew[IS],mut,sel,mu[IS]);

	for (j = 0; j < NM; j++)
		for (i = 0; i < N; i++)
			pop[j][i] = popnew[j][i];
	return;
}

void sample(TYPE* pop, TYPE* popnew, double* mut,double* sel,double mu) {
	double ptot;
	double select[2*N];
	int i,index;
	// sum of probabilities for normalization
	ptot = 0;
	for (i = 0; i < N; i++)
		ptot += sel[i] + mut[i];

	// evaluate probabilities to be selected
	// (even index) and to be selected+mutated
	// (odd indices)
	select[0] = sel[0] / ptot;
	select[1] = mut[0] / ptot + select[0];
	for (i = 1; i < N; i++) {
		select[2* i ] = sel[i] / ptot + select[2* i - 1];
		select[2* i + 1] = mut[i] / ptot + select[2* i ];
	}
	//get N random numbers in [0,1)
	for (i = 0;	i < N;i++) {
		index = gsl_interp_bsearch(select, gsl_rng_uniform(r), 0, 2*N - 2) + 1; //look up entry in select array

		// if index is odd, mutate this sequence with probability mu, otherwise copy
		popnew[i] = (index % 2 == 1 ? mutate(pop[index / 2], mu) : pop[index / 2]);
	}
	return;
}

inline TYPE mutate (TYPE x, double mu) {
  int i;
  TYPE out=x;
  for (i=0; i<L; i++)
    if (gsl_rng_uniform(r) < mu)
      out^=(ONE << i);
  return out;
}

inline int hd (TYPE x, TYPE y) {
  int dist=0;
  TYPE val=x ^ y;
  while (val) {
    ++dist;
    val &= val-1;
  }
  return dist;
}

inline void outbinary(TYPE num) {
	TYPE tmp = num;
	for(int i=0;i<L;i++) {
		int powr = int(pow(2,(L-i-1)));
		int res = tmp/powr;
		fprintf(stderr,"%i",res);
		tmp-=res*powr;
	}
	fprintf(stderr,"\n");
}
